# Continue-Break Judging Service
# Multi-threaded program judging service for C/C++ programs

import os
import sys
import subprocess
import threading
import time
import signal
import psutil
import tempfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Tuple

# Constants
TIME_LIMIT = 2  # seconds
MEMORY_LIMIT = 256 * 1024 * 1024  # 256 MB in bytes
OUTPUT_LIMIT = 10 * 1024 * 1024  # 10 MB in bytes
MAX_THREADS = 96

# Test case structure
class TestCase:
    def __init__(self, input_data: str, expected_output: str):
        self.input_data = input_data
        self.expected_output = expected_output

def parse_test_data(test_data_file: str) -> List[TestCase]:
    """Parse test data file into list of test cases"""
    test_cases = []
    
    try:
        with open(test_data_file, 'r', encoding='utf-8') as f:
            content = f.read().strip()
            
        if not content:
            return test_cases
            
        # Split by blank lines to separate test cases
        test_blocks = content.split('\n\n')
        
        for block in test_blocks:
            lines = block.strip().split('\n')
            if len(lines) >= 2:
                input_data = '\n'.join(lines[:-1]) + '\n'
                expected_output = lines[-1].strip() + '\n'
                test_cases.append(TestCase(input_data, expected_output))
                
    except Exception as e:
        print(f"Error parsing test data: {e}")
        
    return test_cases

def compile_program(source_file: str, compile_dir: str) -> Tuple[str, str]:
    """Compile the C/C++ program and return executable path and error message"""
    try:
        # Determine file extension
        _, ext = os.path.splitext(source_file)
        if ext not in ['.c', '.cpp']:
            return "", "Unsupported file extension. Use .c or .cpp"
        
        # Create executable name
        exe_name = os.path.splitext(os.path.basename(source_file))[0]
        if ext == '.cpp':
            exe_name += '.exe'
        else:
            exe_name += '.exe'
            
        exe_path = os.path.join(compile_dir, exe_name)
        
        # Compile command
        if ext == '.cpp':
            cmd = ['g++', '-O2', '-std=c++17', source_file, '-o', exe_path]
        else:
            cmd = ['gcc', '-O2', source_file, '-o', exe_path]
            
        # Run compilation
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=30  # 30 second compilation timeout
        )
        
        if result.returncode != 0:
            return "", f"Compilation failed: {result.stderr.strip()}"
            
        return exe_path, ""
        
    except subprocess.TimeoutExpired:
        return "", "Compilation timed out"
    except Exception as e:
        return "", f"Compilation error: {str(e)}"

def run_program_with_limits(exe_path: str, input_data: str, test_case_index: int) -> Tuple[str, str, int, int]:
    """Run program with time and memory limits"""
    try:
        # Create temporary file for input
        with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.in') as f:
            f.write(input_data)
            input_file = f.name
            
        try:
            # Run program with timeout
            start_time = time.time()
            
            # Use subprocess with resource limits
            process = subprocess.Popen(
                [exe_path],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                bufsize=1
            )
            
            # Send input data
            stdout, stderr = process.communicate(input=input_data, timeout=TIME_LIMIT)
            
            end_time = time.time()
            execution_time = end_time - start_time
            
            # Check for memory usage (approximate)
            try:
                p = psutil.Process(process.pid)
                memory_usage = p.memory_info().rss
            except:
                memory_usage = 0
                
            # Check output size
            output_size = len(stdout.encode('utf-8'))
            
            # Check return code
            return_code = process.returncode
            
            # Determine result based on execution
            if execution_time > TIME_LIMIT:
                return "TLE", stdout, return_code, memory_usage
            elif memory_usage > MEMORY_LIMIT:
                return "MLE", stdout, return_code, memory_usage
            elif output_size > OUTPUT_LIMIT:
                return "OLE", stdout, return_code, memory_usage
            elif return_code != 0:
                return "RE", stdout, return_code, memory_usage
            else:
                return "AC", stdout, return_code, memory_usage
                
        finally:
            # Clean up temporary input file
            os.unlink(input_file)
            
    except subprocess.TimeoutExpired:
        return "TLE", "", -1, 0
    except Exception as e:
        return "RE", str(e), -1, 0

def compare_outputs(actual_output: str, expected_output: str) -> bool:
    """Compare actual output with expected output (trimming whitespace)"""
    # Remove trailing whitespace and normalize line endings
    actual = actual_output.rstrip('\n\r')
    expected = expected_output.rstrip('\n\r')
    
    # Compare line by line
    actual_lines = actual.split('\n')
    expected_lines = expected.split('\n')
    
    # Remove empty lines at the end
    while actual_lines and not actual_lines[-1].strip():
        actual_lines.pop()
    while expected_lines and not expected_lines[-1].strip():
        expected_lines.pop()
        
    return actual_lines == expected_lines

def process_test_case(test_case: TestCase, exe_path: str, test_case_index: int) -> str:
    """Process a single test case"""
    try:
        # Run the program
        result, output, return_code, memory_usage = run_program_with_limits(
            exe_path, test_case.input_data, test_case_index
        )
        
        # If it's a compile error or runtime error, return immediately
        if result in ["CE", "RE", "TLE", "MLE", "OLE"]:
            return f"Test {test_case_index + 1}: {result}"
        
        # For AC or WA, compare outputs
        if result == "AC":
            if compare_outputs(output, test_case.expected_output):
                return f"Test {test_case_index + 1}: {result}"
            else:
                return f"Test {test_case_index + 1}: WA"
        else:
            return f"Test {test_case_index + 1}: {result}"
            
    except Exception as e:
        return f"Test {test_case_index + 1}: RE"

def main():
    """Main function"""
    # Check command line arguments
    if len(sys.argv) != 3:
        print("Usage: continue-break <program> <test_data>")
        sys.exit(1)
        
    source_file = sys.argv[1]
    test_data_file = sys.argv[2]
    
    # Check if files exist
    if not os.path.exists(source_file):
        print(f"Error: Source file '{source_file}' not found")
        sys.exit(1)
        
    if not os.path.exists(test_data_file):
        print(f"Error: Test data file '{test_data_file}' not found")
        sys.exit(1)
        
    # Parse test data
    test_cases = parse_test_data(test_data_file)
    if not test_cases:
        print("Error: No test cases found in test data file")
        sys.exit(1)
        
    print(f"Processing {len(test_cases)} test cases...")
    
    # Create temporary directory for compilation
    with tempfile.TemporaryDirectory() as temp_dir:
        # Compile the program
        print("Compiling program...")
        exe_path, compile_error = compile_program(source_file, temp_dir)
        
        if compile_error:
            print(f"Compilation failed: {compile_error}")
            print(f"Test 1: CE")
            return
            
        if not exe_path or not os.path.exists(exe_path):
            print("Compilation failed: No executable created")
            print(f"Test 1: CE")
            return
            
        print("Compilation successful")
        
        # Process test cases with threading
        results = []
        
        # Use ThreadPoolExecutor with max 96 threads
        with ThreadPoolExecutor(max_workers=min(MAX_THREADS, len(test_cases))) as executor:
            # Submit all test cases
            future_to_index = {
                executor.submit(process_test_case, test_case, exe_path, i): i 
                for i, test_case in enumerate(test_cases)
            }
            
            # Collect results as they complete
            for future in as_completed(future_to_index):
                try:
                    result = future.result(timeout=TIME_LIMIT + 1)
                    results.append(result)
                except Exception as e:
                    index = future_to_index[future]
                    results.append(f"Test {index + 1}: RE")
        
        # Print all results
        for result in results:
            print(result)

if __name__ == "__main__":
    main()
